import requests
import json

class BaiduAISearch:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://qianfan.baidubce.com/v2/ai_search/web_search"
    
    def search(self, query, **kwargs):
        """
        执行搜索
        
        Args:
            query: 搜索关键词
            **kwargs: 其他可选参数
        """
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        # 构建请求体 - 修改top_k为3，返回3个结果
        data = {
            "messages": [
                {
                    "content": query,
                    "role": "user"
                }
            ],
            "search_source": "baidu_search_v2",
            "resource_type_filter": kwargs.get('resource_type_filter', [{"type": "web", "top_k": 3}]),
            "search_recency_filter": kwargs.get('search_recency_filter', "year")
        }
        
        # 可选参数
        if 'search_filter' in kwargs:
            data['search_filter'] = kwargs['search_filter']
        
        if 'block_websites' in kwargs:
            data['block_websites'] = kwargs['block_websites']
        
        if 'edition' in kwargs:
            data['edition'] = kwargs['edition']
        
        try:
            response = requests.post(self.base_url, headers=headers, json=data, timeout=30)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            return None

# 使用示例
def baidu_search(search_arg):
    
    API_KEY = "bce-v3/ALTAK-08etyG58vh0N4VrAcpSnz/83f6e8929545afb9983f4b44aa09cea5c18261ed"
    
    # 创建搜索实例
    searcher = BaiduAISearch(API_KEY)
    
    # 基础搜索
    result = searcher.search(search_arg)
    
    if result and 'references' in result:
        res = ''
        print("搜索成功!")
        
        # 获取前3个结果
        references = result['references'][:3]
        
        for i, ref in enumerate(references):
            res += f"{ref['id']}. {ref['title']}\n"
            res += f"   链接: {ref.get('url', 'N/A')}\n"
            
            # 前2个结果返回完整内容
            if i < 3:
                res += f"   完整内容: {ref['content']}\n"
            
            res += "\n"  # 添加空行分隔

    else:
        res = "搜索失败!"
        if result and 'code' in result:
            res += (f"错误码: {result['code']}, 错误信息: {result['message']}")

    return res

